import base64
import time

from flask import Flask, render_template, request, make_response
import cv2
import numpy as np
from yolov5.predict import main
from fastdeploy_test import predict
import threading
import multiprocessing
import queue
import os
os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
app = Flask(__name__)
# app._static_folder = "./static"

# 渲染主页


@app.route('/')
def home():
    return render_template('di.html')


# 从主页跳转到关于页面
@app.route('/about')
def about():
    return render_template('home.html')


@app.route('/upload', methods=['POST'])
def upload():
    file = request.files['file']                  # 获取上传的图片文件
    model = request.form.get('model')    # 获取前端发送的模型选择内容
    print(model)
    # model = 'yolov5'
    img_bytes = file.read()                        # 读取图片数据
    nparr = np.frombuffer(img_bytes, np.uint8)     # 将数据转换为numpy数组
    img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # 将numpy数组转换为BGR格式的OpenCV图片

    if model == '多模型融合':
        s = time.time()

        # with multiprocessing.Pool(processes=3) as pool:
        #    result1 = pool.apply_async(main, args=(img_bgr, )).get()
        #    result2 = pool.apply_async(predict, args=(img_bgr, 'ssd')).get()
        #    result3 = pool.apply_async(predict, args=(img_bgr, 'fast_rcnn')).get()

        result1 = predict(img_bgr, 'yolov5')
        result2 = predict(img_bgr, 'ssd')
        result3 = predict(img_bgr, 'fast_rcnn')
        if result1[1] + result2[1] + result3[1] >= 2:
            detect = result2[0]
        else:
            detect = img_bgr
        e = time.time()
        print('total time: ',e-s)


    else:
        detect, _ = predict(img_bgr, model)

    success, buffer = cv2.imencode('.jpg', detect)
    img_data = base64.b64encode(buffer).decode('ascii')  # 将字节数组转换为Base64编码的字符串
    response = make_response(img_data)     # 创建一个响应对象
    response.headers.set('Content-Type', 'image/jpeg')
    response.headers.set('Content-Disposition', 'attachment', filename='output.jpg')
    return response


if __name__ == '__main__':
    app.run(debug=True)